from matplotlib import pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import os


def plot(dir):
    out_dir  = './imgs/' + dir[dir.find('out/')+4:];
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    fig = plt.figure(figsize=(8, 8), facecolor='w')
    s8 = np.loadtxt(dir+'8.txt')
    # print(s8[1])
    s16 = np.loadtxt(dir+'16.txt')
    s32 = np.loadtxt(dir+'32.txt')
    s64 = np.loadtxt(dir+'64.txt')
    norm = np.loadtxt(dir+'norm.txt')
    # 具体函数方法可用 help(function) 查看，如：help(ax.plot_surface)
    
    ax8 = fig.add_subplot(221,projection = '3d')
    ax8.scatter(s8[0], s8[1], s8[2])
    plt.title('n = 8')
    
    ax16 = fig.add_subplot(222,projection = '3d')
    ax16.scatter(s16[0], s16[1], s16[2],c ='green')
    plt.title('n = 16')
    
    ax32 = fig.add_subplot(223,projection = '3d')
    ax32.scatter(s32[0], s32[1], s32[2],s=5,c='red')
    plt.title('n = 32')
    
    ax64 = fig.add_subplot(224,projection = '3d')
    ax64.scatter(s64[0], s64[1], s64[2],s= 1,c='c')
    plt.title('n = 64')
    # plt.show()
    fig.savefig(out_dir+'solution.pdf')
    
    fig_norm = plt.figure()
    xtics = [8,16,32,64]
    plt.plot(xtics,norm[0],label = '1-norm')
    plt.plot(xtics, norm[1],label='2-norm')
    plt.plot(xtics ,norm[2],label = '$\infinity$-norm')
    plt.legend(loc = 'best')
    plt.xticks(xtics)
    plt.title('Norm')
    # plt.show()
    fig_norm.savefig(out_dir+'norm.pdf')
    
    


plot('./out/IM/')
plot('./out/ID/')
plot('./out/IN/')
plot('./out/RM/')
plot('./out/RD/')
plot('./out/RN/')
